require "Q2DViewer"
require "luaQt"

local nl = require"luaNumeric"
local GLUI = require"luaGLShapes"
local nl = require"luaNumeric"
-- local DataSheet = require("DataSheet")
local FEM2D = require"FEM.FEM2D"

----------------------------------------------------------------------------------------------------------------------------------------
--初始化参数
local VecFEM2D=FEM2D:new()
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--初始化
function VecFEM2D:init ()
	self.omega = 2*pi*self.freq;
	self.k0=self.omega/c0;
	print("k0="..self.k0)
	self.mesh = Geometry.create_mesh_soup();
	self.mesh:is_outline(true)
end
----------------------------------------------------------------------------------------------------------------------------------------


----------------------------------------------------------------------------------------------------------------------------------------
--初始化材料
function VecFEM2D:init_material()
	local ntrigs =self.ntrigs;
	for kk=1,ntrigs do
		local info = self.trigs_infos[kk];
		info:alpha(1)
		info:beta(1)
	end
	
	local nedges =self.nedges;
	for kk=1,nedges do
		local info = self.segs_infos[kk];
		info:gamma(0)
		info:q(0)
	end
end
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--第三类边界条件
function VecFEM2D:set_N_BCs(group_tag,r,q)
	local nsegs = self.nsegs;
	local bcs2 = self.segs_groups_tags[group_tag];
	local segs_infos = self.segs_infos;
	local k0 = self.k0;
	local pts = self.pts;
	print("Num of N seg tags:",bcs2)
	
	for ii=1,#bcs2 do
		tag = bcs2[ii]
		local info = segs_infos[tag];
		--info:gamma(-k0*ci);
		info:gamma(r);
		local pt_tag = info:nodes_tag();
		y1 = pts[ pt_tag[1] ][2];
		y2 = pts[ pt_tag[2] ][2];
		y = (y1+y2)/2.
		if y1 < y2 then
			--info:q(-2*k0*ci);
			info:q(-q);
		else
			info:q(q);
		end
	end
end
----------------------------------------------------------------------------------------------------------------------------------------


----------------------------------------------------------------------------------------------------------------------------------------
--初始化矩阵
function VecFEM2D:init_matrix()
	local ne = self.nedges;
	self.spF =  nl.new_sparse_cmatrix(ne,ne)
	self.spA =  nl.new_sparse_matrix(ne,ne)
	self.spB =  nl.new_sparse_matrix(ne,ne)
	--self.spF =  nl.new_eigen_matrix(ne,ne)
	--self.spA =  nl.new_eigen_matrix(ne,ne)
	--self.spB =  nl.new_eigen_matrix(ne,ne)
	
	self.vecb =  nl.new_eigen_vectorxcd(ne)
	self.vecx =  nl.new_eigen_vectorxcd(ne)
	self.vecb:set_zero();
end
----------------------------------------------------------------------------------------------------------------------------------------


----------------------------------------------------------------------------------------------------------------------------------------
--设置矩阵
function VecFEM2D:set_matrix()
	local k02 = self.k0*self.k0;
	local ne = self.nedges;
	
	function fij(v1,v2,b,c)
		return b[v1]*b[v2]+c[v1]*c[v2];
	end
	
	local spA = self.spA
	local spB = self.spB
	local spF = self.spF
	
	local ntrigs = self.ntrigs;
	local trigs_infos = self.trigs_infos;
	local segs_infos = self.segs_infos;
	for kk=1,ntrigs do
		local info = trigs_infos[kk];
		local pt_tag = info:nodes_tag();
		local seg_tag = info:edges_tag();
		local area = info:area();
		local dirs = info:direction();
		local a = info:a();
		local b = info:b();
		local c = info:c();
		local alpha = info:alpha();
		local beta = info:beta();
		for ii=1,3 do	
			vei1 = ii;
			vei2 = (ii)%3+1;	
			lsii = segs_infos[seg_tag[ii]]:length();	
			for jj=1,3 do
				lsjj = segs_infos[seg_tag[jj]]:length();		
				vej1 = jj;
				vej2 = (jj)%3+1;
				Eij = alpha:real()*lsii*lsjj/area/2;
				Fij  = beta:real()*lsii*lsjj/area/48./2;			
				Fij = Fij*( (1+nl.kronecker_delta(vei1,vej1)) * fij(vei2,vej2,b,c)
					 -(1+nl.kronecker_delta(vei1,vej2)) * fij(vei2,vej1,b,c)
					 -(1+nl.kronecker_delta(vei2,vej1)) * fij(vei1,vej2,b,c)
					 +(1+nl.kronecker_delta(vei2,vej2)) * fij(vei1,vej1,b,c) );
				spA[{seg_tag[ii],seg_tag[jj]}] = spA[{seg_tag[ii],seg_tag[jj]}] +
					(Eij)*dirs[ii]*dirs[jj];
				spB[{seg_tag[ii],seg_tag[jj]}] = spB[{seg_tag[ii],seg_tag[jj]}] +
					(Fij)*dirs[ii]*dirs[jj];
				spF[{seg_tag[ii],seg_tag[jj]}] = spF[{seg_tag[ii],seg_tag[jj]}] +
					(Eij-k02*Fij)*dirs[ii]*dirs[jj];
			end
		end
	end
end
----------------------------------------------------------------------------------------------------------------------------------------


----------------------------------------------------------------------------------------------------------------------------------------
--强加第一类边界条件
function VecFEM2D:set_PEC_BCs(group_tag,array_M,array_b)
	local nedges = self.nedges;
	local bcs2 = self.segs_groups_tags[group_tag];
	for ii=1,#bcs2 do
		tag = bcs2[ii];
		for jj=1,nedges do
			for k,spM in pairs(array_M) do
				if nl.abs(spM[{tag,jj}]) >1e-20 then
					spM[{tag,jj}] = 0
				end
				if nl.abs(spM[{jj,tag}]) >1e-20 then
					spM[{jj,tag}] = 0
				end
			end
		end
		for k,spM in pairs(array_M) do
			spM[{tag,tag}] = 100
		end
		for k,b in pairs(array_b) do
			--b[tag] =  0
		end
	end
end
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--第三类边界条件
function VecFEM2D:set_N_BCs_matrix()
	local nedges = self.nedges;
	local segs_infos = self.segs_infos;	
	for kk=1,nedges do
		local info = segs_infos[kk];
		local pt_tag = info:nodes_tag();
		local ls = info:length();
		local q = info:q();
		local r = info:gamma();
		if nl.abs(r) > 1e-5 then
			self.spF[{kk,kk}] = self.spF[{kk,kk}]+ 0.5*r*ls;
			self.vecb[kk] = self.vecb[kk] + q*ls;
		end
	end
end
----------------------------------------------------------------------------------------------------------------------------------------


----------------------------------------------------------------------------------------------------------------------------------------
--求解
function VecFEM2D:solve()
	local ts = os.clock()
	--assert( spm:solve_ldlt (vecb,vecx)>=0 , "solve failed")
	--assert( spm:solve_cg (vecb,vecx)>=0 , "solve failed")
	assert( self.spF:solve_lu(self.vecb,self.vecx)>=0 , "solve failed")
	print(string.format("elapsed time: %.2f\n", os.clock() - ts))
end
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--后处理
function VecFEM2D:post_calculate()
	local ntrigs = self.ntrigs;
	local fields = self.mesh:points_vals();
	local trigs_infos = self.trigs_infos;
	local segs_infos = self.segs_infos;
	local mesh = self.mesh;
	local vecx = self.vecx;
	local k0 = self.k0;
	
	
	mesh:show_facet_fields();
	fields:resize( ntrigs )
	
	local cfields =  nl.new_eigen_vectorxcd(ntrigs)

	for kk=1,ntrigs do
		local info = trigs_infos[kk];
		local seg_tag = info:edges_tag();
		local area = info:area();
		local dirs = info:direction();
		local info = segs_infos[kk];

		cfields[kk] = 0;
		for jj=1,3 do
			local seg_info = segs_infos[seg_tag[jj]];
			cfields[kk] = cfields[kk]  + seg_info:length() * vecx[seg_tag[jj]] * dirs[jj];
		end
		cfields[kk] = cfields[kk] /area; 
	end
	for ii=1,ntrigs do
		fields[ii] = cfields[ii]:real();
	end
	Geometry.AddShape(mesh);
end
----------------------------------------------------------------------------------------------------------------------------------------

return VecFEM2D;